﻿/********************************/
/****** TCP Socket FOR WIN  *****/
/********************************/
//SOCKET.cpp
#include "stdafx.h"
#include "stdio.h"
#include <sys/types.h>
#include <winsock.h>
#include <time.h>
#include "string.h"
#include "signal.h"
#include "errno.h"
#include "socket.h"


///////////////////////////////////////////

CSock::CSock(int socket,CSocket* pSocket)
{
    m_socket = socket;
    m_pParent = pSocket;
    
    int ntimev = 0;
    setsockopt(socket, SOL_SOCKET, SO_RCVTIMEO, (char*)&ntimev, sizeof(ntimev));
    setsockopt(socket, SOL_SOCKET, SO_SNDTIMEO, (char*)&ntimev, sizeof(ntimev));

    m_sendBuffer = new unsigned char[SIZE];
    m_recvBuffer = new unsigned char[SIZE];
    m_sendStart = m_sendEnd = 0;               //Initialize  start and end flag
    m_recvStart = m_recvEnd = 0;

    InitializeCriticalSection(&m_cs);
    m_event1 = CreateEvent(NULL,false,false,NULL);
    m_event2 = CreateEvent(NULL,false,false,NULL);

    DWORD threadid;
    m_sendThread = CreateThread(NULL,0,SendThread,this,NULL,&threadid);
    m_recvThread = CreateThread(NULL,0,RecvThread,this,NULL,&threadid);
    m_watchThread = CreateThread(NULL,0,WatchThread,this,NULL,&threadid);
}

///////////////////////////////////////////////

CSock::~CSock()
{
    if(0 != m_watchThread)
    {
        TerminateThread(m_watchThread,0);
        CloseHandle(m_watchThread);
    }

    if(0 != m_sendThread)
    {
        TerminateThread(m_sendThread,0);
        CloseHandle(m_sendThread);
    }
    
    if(0 != m_recvThread)
    {
        TerminateThread(m_recvThread,0);
        CloseHandle(m_recvThread);
    }

    m_pParent->EndSocket(this);

    if(m_socket != 0)
     {
         shutdown(m_socket,2);
         closesocket(m_socket);
     } 	
    
    delete m_sendBuffer;
    delete m_recvBuffer;
    
    DeleteCriticalSection(&m_cs);
    CloseHandle(m_event1);
    CloseHandle(m_event2);
//	m_sendThread = 0;
//	m_recvThread = 0;
}

////////////////////////////////////////

int CSock::Send(char* buffer,int len)
{
    if(len == 0) return 0; 
    EnterCriticalSection(&m_cs);
    
    int valid = 0;
    if(m_sendEnd >= m_sendStart)
        valid = SIZE-(m_sendEnd-m_sendStart)-1;
    else
        valid = m_sendStart-m_sendEnd-1;
        
    if(valid < len+8)
    {
        LeaveCriticalSection(&m_cs);
        return 0;
    }

    unsigned int a[2];
    a[0] = 0x55aa55aa;
    a[1] = htonl(len);
    
    for(int i=0;i<8;i++)
    {
        m_sendBuffer[m_sendEnd] = ((char*)a)[i];
        m_sendEnd = (m_sendEnd+1) % SIZE;
    }
    
    for(int i=0;i<len;i++)
    {
        m_sendBuffer[m_sendEnd] = buffer[i];
        m_sendEnd = (m_sendEnd+1) % SIZE;
    }

    SetEvent(m_event1);
    LeaveCriticalSection(&m_cs);

    return len;	
}
////////////////////////////////////////////////

int CSock::CheckValidData()          
{
    EnterCriticalSection(&m_cs);
        
    if(m_recvStart == m_recvEnd)      //at the beginning of sending data 
    {
        LeaveCriticalSection(&m_cs);
        return 0;
    }
        
    int valid = 0;                          //valid length in recv buffer
    if(m_recvEnd >= m_recvStart)
        valid = m_recvEnd-m_recvStart;
    else
        valid = SIZE-(m_recvStart-m_recvEnd);
    
    if(valid < 8)                          //no valid data
    {
        LeaveCriticalSection(&m_cs);
        return 0;
    }
    else    
    {
        //check data packet head
        if(m_recvBuffer[m_recvStart] == 0xaa && m_recvBuffer[(m_recvStart+1)%SIZE] == 0x55 &&
            m_recvBuffer[(m_recvStart+2)%SIZE] == 0xaa && m_recvBuffer[(m_recvStart+3)%SIZE] == 0x55)
        {
            INT32* pNetLen = (INT32*) (m_recvBuffer + m_recvStart + 4);

            INT32 length = ntohl(*pNetLen);
           
            if(length > SIZE)    
            {
                m_recvStart = m_recvEnd = 0;
                LeaveCriticalSection(&m_cs);
                SetEvent(m_event2);
                return 0;
            }
            else
            if(valid < 8+(int)length)     // no valid data
            {
                LeaveCriticalSection(&m_cs);
                return 0;		
            }
            
            if(length == 16) // tick data ?
            {
                int td;
                for(int i=0;i<4;i++)
                    ((unsigned char*)&td)[i] = m_recvBuffer[(m_recvStart+8+i)%SIZE];
                
                if(td == 0xaaaaaaaa) // ignore tick data
                {
                    m_recvStart = (m_recvStart+8+length)%SIZE;
                    length = 0;
                }
            }
            
            LeaveCriticalSection(&m_cs);
            SetEvent(m_event2);
            
            return length;
        }
        else   // when data packet head is wrong 
        {
            m_recvStart = m_recvEnd = 0;
            LeaveCriticalSection(&m_cs);
            SetEvent(m_event2);
            return 0;
        }		
    }
    
    LeaveCriticalSection(&m_cs);
    return 0;
}

///////////////////////////////////////////////

int CSock::Recv(char* buffer,int len)
{
    EnterCriticalSection(&m_cs);
    
    if(m_recvStart == m_recvEnd)
    {
        LeaveCriticalSection(&m_cs);
        return 0;
    }
        
    int valid = 0;
    if(m_recvEnd >= m_recvStart)
        valid = m_recvEnd-m_recvStart;
    else
        valid = SIZE-(m_recvStart-m_recvEnd);
    
    if(valid < 8)
    {
        LeaveCriticalSection(&m_cs);
        return 0;
    }
    else
    {
        if(m_recvBuffer[m_recvStart] == 0xaa && m_recvBuffer[(m_recvStart+1)%SIZE] == 0x55 &&
            m_recvBuffer[(m_recvStart+2)%SIZE] == 0xaa && m_recvBuffer[(m_recvStart+3)%SIZE] == 0x55)
        {
           
            unsigned int netLength = 0;
            char* p = (char*)&netLength;
            p[0] = m_recvBuffer[(m_recvStart+4)%SIZE];
            p[1] = m_recvBuffer[(m_recvStart+5)%SIZE];
            p[2] = m_recvBuffer[(m_recvStart+6)%SIZE];
            p[3] = m_recvBuffer[(m_recvStart+7)%SIZE];
            
            unsigned int length = ntohl(netLength);
            
            if(length > SIZE) 
            {
                m_recvStart = m_recvEnd = 0;
                LeaveCriticalSection(&m_cs);
                SetEvent(m_event2);
                return 0;
            }
            else
            if((DWORD)valid < 8+length)
            {
                LeaveCriticalSection(&m_cs);
                return 0;		
            }
            
            if(length > (DWORD)len)
            {
                LeaveCriticalSection(&m_cs);
                return -1;
            }
            else
            {
                for(unsigned int i=0;i<length;i++)
                {
                    buffer[i] = m_recvBuffer[(m_recvStart+8+i)%SIZE];
                }
                m_recvStart = (m_recvStart+8+length)%SIZE;
                LeaveCriticalSection(&m_cs);
                SetEvent(m_event2);				

                if(length == 16)
                {
                    DWORD* p = (DWORD*)buffer;
                    if(p[0] == 0xaaaaaaaa)
                    {
                        //AfxMessageBox("tick\r\n");
                        return 0;
                    }
                
                }
                
                return length;
            }
        }
        else
        {
            m_recvStart = m_recvEnd = 0;
            LeaveCriticalSection(&m_cs);
            SetEvent(m_event2);
            return 0;
        }		
    }
    
    LeaveCriticalSection(&m_cs);
    return 0;
}

///////////////////////////////////////

DWORD CSock::SendThread(void* arg)
{
    CSock* pThis = (CSock*)arg;
    
    while(1)
    {
        int valid = 0;

        EnterCriticalSection(&pThis->m_cs);
        
        if(pThis->m_sendEnd >= pThis->m_sendStart)
            valid = pThis->m_sendEnd-pThis->m_sendStart;
        else
            valid = SIZE-(pThis->m_sendStart-pThis->m_sendEnd);
            
        int sendLen = 0;
        if(valid >= 1*1024) sendLen = 1*1024;
        else sendLen = valid;
    
        if(sendLen >  0)
        {
            if(pThis->m_sendStart + sendLen > SIZE)
                sendLen = SIZE-pThis->m_sendStart;
            
            unsigned char tempBuf[1*1024];
            
            for(int i=0;i<sendLen;i++)
                tempBuf[i] = pThis->m_sendBuffer[pThis->m_sendStart+i];
                    
            LeaveCriticalSection(&pThis->m_cs);
            
            int length = send(pThis->m_socket,(char*)tempBuf,sendLen,0);
            
            if(length >= 0)     // modify start flag
            {
                pThis->m_sendStart = (pThis->m_sendStart+length) % SIZE;
            }
            else   // send failed, close socket
            {
                HANDLE h = pThis->m_sendThread;
                pThis->m_sendThread = 0;
                delete pThis;
                CloseHandle(h);
                printf("Send thread close socket");	
                return 0;
            }			
        }
        else
        {
            LeaveCriticalSection(&pThis->m_cs);
            ResetEvent(pThis->m_event1);
            WaitForSingleObject(pThis->m_event1,INFINITE);
        }
    }
}

DWORD CSock::WatchThread(void* arg)
{
    
    CSock* pThis = (CSock*)arg;
    
    DWORD Buf[4] = {0xaaaaaaaa,0xaaaaaaaa,0xaaaaaaaa,0xaaaaaaaa};
    
    int watchFlag = pThis->m_watchFlag;

    pThis->m_nCount = 0;
    while(1)
    {
        Sleep(500);
        pThis->Send((char*)Buf,16);

        pThis->m_nCount++;
        if((pThis->m_nCount % 35) == 0)
        {
            if(watchFlag == pThis->m_watchFlag)
            {
                HANDLE h = pThis->m_watchThread;
                pThis->m_watchThread = 0;
                delete pThis;
                CloseHandle(h);
                printf("Watch thread close socket");	
                return 0;
            }

            watchFlag = pThis->m_watchFlag;
        }
    }
}

//////////////////////////////////////////////

DWORD CSock::RecvThread(void* arg)
{
    CSock* pThis = (CSock*)arg;
    
    char tempBuf[4*1024];
    
    while(1)
    {
        int len = recv(pThis->m_socket,tempBuf,4*1024,0);
        if(len <= 0) 
        {
            HANDLE h = pThis->m_recvThread;
            pThis->m_recvThread = 0;
            delete pThis;
            CloseHandle(h);
            printf("Recv thread close socket");	
            return 0;
        }
        else
        {
            pThis->m_watchFlag ++;
            pThis->m_nCount=0;
            while(1)
            {
                EnterCriticalSection(&pThis->m_cs);	
                
                int valid = 0;
                if(pThis->m_recvEnd >= pThis->m_recvStart)
                    valid = SIZE-(pThis->m_recvEnd-pThis->m_recvStart)-1;
                else
                    valid = pThis->m_recvStart-pThis->m_recvEnd-1;
                
                if(valid > len) break;
    
                LeaveCriticalSection(&pThis->m_cs);
            
                ResetEvent(pThis->m_event2);
                WaitForSingleObject(pThis->m_event2,0);
            }

            for(int i=0;i<len;i++)
            {
                pThis->m_recvBuffer[(pThis->m_recvEnd+i)%SIZE] = tempBuf[i];
            }
            
            pThis->m_recvEnd = (pThis->m_recvEnd+len)%SIZE;
            
            LeaveCriticalSection(&pThis->m_cs);
            
            pThis->m_pParent->OnRecv(pThis);
        }
    }
}

///////////////////////////////////////////




///////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////


void CSocket::OnRecv(CSock* pSocket)
{
    while(1)          
    {
        if(m_pRecvCallBack != NULL && pSocket->CheckValidData() > 0)   //check received data 
            m_pRecvCallBack(pSocket,m_msg);   // process data packet one by one 
        else
            break;      // no callback or valid data = 0 when data processing finished 
    }
}

////////////////////////////////////////////////////

void CSocket::EndSocket(CSock* pSocket)
{
    int pos = -1;
    EnterCriticalSection(&m_cs);
    
    for(int i=0;i<m_lineCount;i++)
    {
        if(m_pSocket[i] == pSocket)
        {
            pos = i;break;
        }
    }
    
    if(pos >= 0)
    {
        for(int i=pos;i<m_lineCount-1;i++)
        {
            m_pSocket[i] = m_pSocket[i+1];
        }
        m_lineCount--;
    }
    LeaveCriticalSection(&m_cs);
}

///////////////////////////////////////////////

CSocket::CSocket(int port, bool bServer)
{
    static bool inited = false;

    InitializeCriticalSection(&m_cs);
    
    if(!inited)
    {
        inited = true;
    }
    
    m_bServer = bServer;
    m_port = port;

    m_lineCount = 0;
    m_pRecvCallBack = NULL;

    m_listenSocket = 0;
    m_listenThread = 0;
}

////////////////////////////////////////

CSocket::~CSocket()
{
    while(GetValidLineCount() > 0)
    {
        delete GetLineByIndex(0);
    }

    if(0 != m_listenThread)
    {
        TerminateThread(m_listenThread,0);
        CloseHandle(m_listenThread);
    } 

    if(m_listenSocket != 0) 
    {
        shutdown(m_listenSocket,2);
        closesocket(m_listenSocket);
    }

    DeleteCriticalSection(&m_cs);
}

////////////////////////////////////////////

bool CSocket::Listen(int maxConnectCount)
{
    m_maxConnectCount = maxConnectCount;
    
    m_listenSocket = socket(AF_INET,SOCK_STREAM, IPPROTO_TCP);
    
    if(m_listenSocket == -1) return false;
    
    sockaddr_in addr;
    memset(&addr,0,sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(m_port);
    addr.sin_addr.s_addr = INADDR_ANY;
    bind(m_listenSocket,(sockaddr*)&addr,sizeof(addr));
    
    if(listen(m_listenSocket,1) == -1) return false;

    DWORD threadid;
    m_listenThread = CreateThread(NULL,0,ListenThread,this,0,&threadid);

    return true;
}

//////////////////////////////////////////////

DWORD CSocket::ListenThread(void* arg)
{
    CSocket* pSocket = (CSocket*)arg;
    
    while(1)

    {
        sockaddr_in ra;
        size_t len = sizeof(ra);
        int rsock = accept(pSocket->m_listenSocket,(sockaddr*)&ra,(int*)&len);
        if(rsock != -1 && rsock != 0)
        {
            EnterCriticalSection(&pSocket->m_cs);
            
            if(pSocket->GetValidLineCount() >= pSocket->m_maxConnectCount)
            {
                shutdown(rsock,2);
                closesocket(rsock);
            }
            else
            {
                pSocket->m_pSocket[pSocket->m_lineCount] = new CSock(rsock,pSocket);
                pSocket->m_lineCount++;
            }
            LeaveCriticalSection(&pSocket->m_cs);
        }
    }
}	

//////////////////////////////////////

bool CSocket::Connect(char ipAddr[],int port)
{
    int sock = socket(AF_INET,SOCK_STREAM,0);
    
    sockaddr_in addr;
    memset(&addr,0,sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = INADDR_ANY;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    bind(sock,(sockaddr*)&addr,sizeof(addr));
    
    memset(&addr,0,sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.S_un.S_addr = inet_addr(ipAddr);
    
    if(connect(sock,(sockaddr*)&addr,sizeof(addr)) == 0)
    {
        EnterCriticalSection(&m_cs);
        m_pSocket[m_lineCount] = new CSock(sock,this);
        m_lineCount++;
        LeaveCriticalSection(&m_cs);
        return true;
    }
    else
    {
        shutdown(sock,2);
        closesocket(sock);
        return false;
    }
    return true;
}

//////////////////////////////////////////////	

int CSocket::GetValidLineCount()
{
    return m_lineCount;
}

/////////////////////////////////////////////

CSock* CSocket::GetLineByIndex(int nIndex)
{
    if(nIndex > m_lineCount-1) return NULL;
    return m_pSocket[nIndex];
}

/////////////////////////////////////////////

int CSocket::SendData(CSock* pSocket,char* buffer,int len)
{
    for(int i=0;i<m_lineCount;i++)
    {
        if(m_pSocket[i] == pSocket)
        {
            return pSocket->Send(buffer,len);
        }
    }
    return 0;
}

/////////////////////////////////////////

int CSocket::RecvData(CSock* pSocket,char* buffer,int len)
{
    for(int i=0;i<m_lineCount;i++)
    {
        if(m_pSocket[i] == pSocket)
        {
            return pSocket->Recv(buffer,len);
        }
    }
    return 0;
}

/////////////////////////////////	
    
int CSocket::SetCallBack(void (*proc)(CSock* pSocket,int MSG))
{
    m_pRecvCallBack = proc;
    return 0;
}

//////////////////////////////////

void CSocket::SetMsg(int msg)
{
    m_msg = msg;
}
////////////////////////////